import os

def get_image_paths(folder, extensions):
    image_paths = []
    for root, dirs, files in os.walk(folder):
        for file in files:
            if file.lower().endswith(extensions):
                image_paths.append(os.path.join(root, file))
    return image_paths

def save_paths_to_txt(image_paths, output_file):
    with open(output_file, 'w') as f:
        for path in image_paths:
            f.write(path + '\n')

# 设置要遍历的文件夹路径列表
# folders = ['/path/to/folder1', '/path/to/folder2']
# folders = ['/home/yjh_working/data/helmet_data', '/home/yjh_working/data/public_helmet_data/21-SHWD','/home/yjh_working/data/public_helmet_data/22-SFCHD']
folders = ['/home/yjh_test/data/helmet_data/HelmetDataset-zhengshu-20241121']
# 设置图片文件的扩展名
extensions = ('.jpg', '.jpeg', '.png', '.gif', '.bmp')
# extensions = ('.txt')
# 设置输出文本文件的路径
output_file = 'all_image.txt'

# 初始化图片路径列表
all_image_paths = []

# 遍历每个文件夹并获取图片路径
for folder in folders:
    all_image_paths.extend(get_image_paths(folder, extensions))


# 将图片路径写入文本文件
save_paths_to_txt(all_image_paths, output_file)

print(f"Image paths have been written to {output_file}")